#!/usr/bin/env python
# Import the os module to interact with the operating system, primarily for accessing environment variables.
import os


# This function checks if a predefined list of tables exists in a specific database schema.
# @pysnooper.snoop()
def check_tables():
    # Retrieve the database connection string from an environment variable named 'MYSQL_SERVICE'.
    SQLALCHEMY_DATABASE_URI = os.getenv('MYSQL_SERVICE', '')
    # Proceed only if the environment variable is set.
    if SQLALCHEMY_DATABASE_URI:
        # Import the 'url' class from sqlalchemy.engine to parse the database URI.
        import sqlalchemy.engine.url as url

        # Parse the connection string into a URL object to easily access its components (host, port, etc.).
        uri = url.make_url(SQLALCHEMY_DATABASE_URI)
        """Inits the Myapp application"""
        # Import the pymysql library, which is a pure-Python MySQL client.
        import pymysql

        # 创建连接
        # Establish a connection to the MySQL database using the parsed credentials.
        conn = pymysql.connect(
            host=uri.host, port=uri.port, user=uri.username, password=uri.password, charset='utf8'
        )
        # 创建游标
        # Create a cursor object from the connection, which allows executing SQL queries.
        cursor = conn.cursor()

        # 创建数据库的sql(如果数据库存在就不创建，防止异常)
        # Define an SQL query to select all table names from the 'information_schema.tables' view
        # where the table schema is 'kubeflow'.
        sql = "SELECT table_name FROM information_schema.tables  WHERE table_schema='kubeflow'"
        # Execute the SQL query.
        cursor.execute(sql)
        # Fetch all the rows from the query result.
        results = list(cursor.fetchall())
        # The result is a list of tuples, so we extract the first element (the table name) from each tuple.
        results = [item[0] for item in results]
        # Print the list of existing tables for debugging purposes.
        print(results)
        # Define a list of table names that are expected to exist in the 'kubeflow' schema.
        for table_name in [
            'ab_permission',
            'ab_permission_view',
            'ab_permission_view_role',
            'ab_register_user',
            'ab_role',
            'ab_user',
            'ab_user_role',
            'ab_view_menu',
            'alembic_version',
            'dimension',
            'docker',
            'images',
            'inferenceservice',
            'job_template',
            'logs',
            'metadata_metric',
            'model',
            'nni',
            'notebook',
            'pipeline',
            'project',
            'project_user',
            'pytorchjob',
            'repository',
            'run',
            'service',
            'service_pipeline',
            'task',
            'metadata_table',
            'tfjob',
            'user_attribute',
            'workflow',
            'xgbjob',
        ]:
            # Check if an expected table is missing from the list of tables fetched from the database.
            if table_name not in results:
                # If a table is missing, print an error message with instructions.
                print(
                    'kubeflow db下，table %s不完整，请\n1、drop dabatase kubeflow\n2、重启当前pod'
                    % table_name
                )
                # Exit the script with a non-zero status code to indicate failure.
                exit(1)


# Call the function to execute the check.
check_tables()
